import torch
import torch.nn.functional as F

import re
# from torch._six import container_abcs, string_classes, int_classes


np_str_obj_array_pattern = re.compile(r'[SaUO]')


default_collate_err_msg_format = (
    "default_collate: batch must contain tensors, numpy arrays, numbers, "
    "dicts or lists; found {}")


def collate_function(batch):
    r"""Puts each data field into a tensor with outer dimension batch size"""
    # print("collate:",len(batch))
    images = []
    labels  = []
    for item in batch:
        # print("image:",item)
        # print('batch',batch)
        images.append(item['img'])
        labels.append(item['label'])
    return (torch.stack(images,dim=0),torch.stack(labels,dim=0))
